{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from math import ceil\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "from utils.input_pipeline import get_image_folders\n",
    "from utils.training import train, optimization_step\n",
    "from utils.diagnostic import count_params\n",
    "    \n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create data iterators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_size = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100000"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_folder, val_folder = get_image_folders()\n",
    "\n",
    "train_iterator = DataLoader(\n",
    "    train_folder, batch_size=batch_size, num_workers=4,\n",
    "    shuffle=True, pin_memory=True\n",
    ")\n",
    "\n",
    "val_iterator = DataLoader(\n",
    "    val_folder, batch_size=256, num_workers=4,\n",
    "    shuffle=False, pin_memory=True\n",
    ")\n",
    "\n",
    "# number of training samples\n",
    "train_size = len(train_folder.imgs)\n",
    "train_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# number of validation samples\n",
    "val_size = len(val_folder.imgs)\n",
    "val_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from get_densenet import get_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model, loss, optimizer = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7151176"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# number of params in the model\n",
    "count_params(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "391"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_epochs = 5\n",
    "n_batches = ceil(train_size/batch_size)\n",
    "\n",
    "# total number of batches in the train set\n",
    "n_batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0  3.882 2.772  0.208 0.392  0.427 0.659  776.920\n",
      "1  3.110 2.502  0.308 0.425  0.565 0.690  761.085\n",
      "2  2.974 2.403  0.330 0.441  0.588 0.704  761.199\n",
      "3  2.925 2.367  0.338 0.444  0.596 0.706  761.763\n",
      "4  2.891 2.329  0.344 0.455  0.603 0.719  760.764\n",
      "CPU times: user 50min 25s, sys: 13min 12s, total: 1h 3min 38s\n",
      "Wall time: 1h 3min 41s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "def optimization_step_fn(model, loss, x_batch, y_batch):\n",
    "    return optimization_step(model, loss, x_batch, y_batch, optimizer)\n",
    "\n",
    "all_losses = train(\n",
    "    model, loss, optimization_step_fn,\n",
    "    train_iterator, val_iterator, n_epochs\n",
    ")\n",
    "# epoch logloss  accuracy    top5_accuracy time  (first value: train, second value: val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "model.cpu();\n",
    "torch.save(model.state_dict(), 'model_step1.pytorch_state')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
